import os
from dotenv import load_dotenv

from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent

from my_deep_agents_from_scratch.prompts import WRITE_TODOS_DESCRIPTION
from my_deep_agents_from_scratch.prompts import TODO_USAGE_INSTRUCTIONS
from my_deep_agents_from_scratch.state import DeepAgentState
from my_deep_agents_from_scratch.todo_tools import read_todos, write_todos
from utils import format_messages, show_prompt

def main():
    load_dotenv(os.path.join("..", ".env"), override=True)
    model = os.getenv("MODEL")
    base_url = os.getenv("BASE_URL")
    api_key = os.getenv("_API_KEY")

    # show_prompt(WRITE_TODOS_DESCRIPTION)
    # show_prompt(TODO_USAGE_INSTRUCTIONS)

    # 模拟 search result
    search_result = """模型上下文协议（MCP）是由Anthropic开发的一种标准协议，旨在实现AI模型与外部系统（如工具、数据库及其他服务）的无缝集成。
    它作为标准化的通信层，使AI模型能够以一致高效的方式访问并利用来自不同数据源的信息。
    本质上，MCP通过为数据交换提供统一语言，极大简化了将AI助手连接到外部服务的过程。"""

    # 模拟 search tool
    @tool(parse_docstring=True)
    def web_search(
            query: str,
    ):
        """Search the web for information on a specific topic.

        This tool performs web searches and returns relevant results
        for the given query. Use this when you need to gather information from
        the internet about any topic.

        Args:
            query: The search query string. Be specific and clear about what
                   information you're looking for.

        Returns:
            Search results from search engine.

        Example:
            web_search("machine learning applications in healthcare")
        """
        return search_result

    # Create agent using create_react_agent directly
    model = ChatOpenAI(model=model, base_url=base_url, api_key=api_key)
    tools = [write_todos, web_search, read_todos]

    # Add mock research instructions
    SIMPLE_RESEARCH_INSTRUCTIONS = """IMPORTANT: Just make a single call to the web_search tool and use the result provided by the tool to answer the user's question."""

    # Create agent
    agent = create_react_agent(
        model,
        tools,
        prompt=TODO_USAGE_INSTRUCTIONS
               + "\n\n"
               + "=" * 80
               + "\n\n"
               + SIMPLE_RESEARCH_INSTRUCTIONS,
        state_schema=DeepAgentState,
    )

    # Show the agent
    png_data = agent.get_graph(xray=True).draw_mermaid_png()
    output_path = "to_agent_graph.png"  # 保存到文件
    with open(output_path, "wb") as f:
        f.write(png_data)

    # Example usage "从初始状态开始，待办事项列表为空，并收到一个用户的research request。"
    result = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "请给我一份模型上下文协议（MCP）的简要概述。",
                }
            ],
            "todos": []
        }
    )

    format_messages(result["messages"])

if __name__ == '__main__':
    main()